package com.shujia.core

import org.apache.spark.rdd.RDD
import org.apache.spark.{SparkConf, SparkContext}

object Demo13AggregateByKey {
  def main(args: Array[String]): Unit = {
    // aggregateByKey：转换算子，按照指定的计算逻辑，对Key分组后进行预聚合以及聚合操作
    val conf: SparkConf = new SparkConf()
    conf.setMaster("local")
    conf.setAppName("Demo13AggregateByKey")

    val sc: SparkContext = new SparkContext(conf)
    val stuRDD: RDD[String] = sc
      .textFile("spark/data/stu/students.txt")

    // 统计班级人数
    val stuKVRDD: RDD[(String, Int)] = stuRDD
      // 将数据变成KV格式，以班级作为Key，1作为Value
      .map(line => {
        val splits: Array[String] = line.split(",")
        (splits(4), 1)
      })

    // 将数据变成KV格式，以班级作为Key，age作为Value
    val stuAgeRDD: RDD[(String, Int)] = stuRDD.map(line => {
      val splits: Array[String] = line.split(",")
      (splits(4), splits(2).toInt)
    })

    stuKVRDD.foreach(println)

    stuKVRDD.reduceByKey(_ + _).foreach(println)

    /**
     * aggregateByKey接收两组参数：
     * 第一组参数zeroValue：初始化的值，类型基于值可以推断的，可以指定任意的值
     * 第二组参数需要接收两个函数：seqOp、combOp
     * seqOp：预聚合函数，作用在每个MapTask内部
     * combOp：聚合函数，作用在每个ReduceTask内部
     */
    stuKVRDD.aggregateByKey(0)((i1, i2) => {
      i1 + i2
    }, (u1, u2) => {
      u1 + u2
    }).foreach(println)

    // 统计班级平均年龄

    // 先统计班级人数
    val clazzCntRDD: RDD[(String, Int)] = stuKVRDD.reduceByKey(_ + _)
    // 统计班级年龄之和
    val classAgeRDD: RDD[(String, Int)] = stuAgeRDD.reduceByKey(_ + _)

    // 通过班级进行关联
    clazzCntRDD.join(classAgeRDD).map {
      case (clazz: String, (cnt: Int, sumAge: Int)) =>
        s"$clazz,${sumAge / cnt.toDouble}"
    }.foreach(println)

    // 通过aggregateByKey实现
    stuAgeRDD.aggregateByKey(0)((sumAge, age) => {
      sumAge + age
    }, (sumAge1, sumAge2) => {
      sumAge1 + sumAge2
    }).foreach(println)

    stuKVRDD.aggregateByKey(0)((i1, i2) => {
      i1 + 1
    }, (u1, u2) => {
      u1 + u2
    }).foreach(println)

    stuAgeRDD.aggregateByKey((0, 0))((t2, age) => {
      // 对age进行预聚合
      val sumAge: Int = t2._1 + age
      // 对班级人数+1
      val clazzCnt: Int = t2._2 + 1
      (sumAge, clazzCnt)
    }, (sumAgeAndClazzCnt01, sumAgeAndClazzCnt02) => {
      val totalSumAge: Int = sumAgeAndClazzCnt01._1 + sumAgeAndClazzCnt02._1
      val totalClazzCnt: Int = sumAgeAndClazzCnt01._2 + sumAgeAndClazzCnt02._2
      (totalSumAge, totalClazzCnt)
    }).map {
      case (clazz: String, (sumAge: Int, cnt: Int)) =>
        s"$clazz,${sumAge / cnt.toDouble}"
    }.foreach(println)

    // 使用reduceByKey简化aggregateByKey
    // 需要先将学生数据转换成KV格式
    // 用clazz作为key，用(age,1)构建二元组作为Value
    stuRDD.map(line => {
      val splits: Array[String] = line.split(",")
      (splits(4), (splits(2).toInt, 1))
    }).reduceByKey((t1, t2) => {
      val sumAge: Int = t1._1 + t2._1
      val cnt: Int = t1._2 + t2._2
      (sumAge, cnt)
    }).map {
      case (clazz: String, (sumAge: Int, cnt: Int)) =>
        s"$clazz,${sumAge / cnt.toDouble}"
    }.foreach(println)

  }

}
